package com.rpj.stauy.controller;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import io.qdrant.client.grpc.Collections;
import jakarta.annotation.Resource;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;


@RestController
@RequestMapping
public class EmbeddinglController {

    //原始向量model 等价ChatModel
    @Resource
    private EmbeddingModel embeddingModel;

    //向量数据库客户端
    @Resource
    private QdrantClient qdrantClient;

    @Resource
    private EmbeddingStore<TextSegment> embeddingStore;

    //查看生成的向量数据
    //http://localhost:9012/embedding/embed
    @GetMapping("/embedding/text")
    public String embeddingText(){
        String prompt = """
                   咏鸡
                鸡鸣破晓光，
                红冠映朝阳。
                金羽披霞彩，
                昂首步高岗。
                """;
        //内容转为向量数据
        Response<Embedding> embeddingResponse = embeddingModel.embed(prompt);
        System.out.println(embeddingResponse);
        return embeddingResponse.content().toString();
    }


    //新建向量数据库实例和创建索引
    //http://localhost:9012/embedding/createCollection
    @GetMapping("/embedding/createCollection")
    public void createCollection(){
        Collections.VectorParams vectorParams = Collections.VectorParams.newBuilder()
                .setDistance(Collections.Distance.Cosine)
                .setSize(1024)
                .build();
        qdrantClient.createCollectionAsync("test-qdrant",vectorParams);

    }


    //向向量数据库中添加文本

    //http://localhost:9012/embedding/add
    @GetMapping(value = "/embedding/add")
    public String add(){

        String prompt = """
                   咏鸡
                鸡鸣破晓光，
                红冠映朝阳。
                金羽披霞彩，
                昂首步高岗。
                """;
        //包装文本
        TextSegment textSegment = TextSegment.from(prompt);
        //给包装后的对象添加元数据，用于存储与文本相关的额外信息
        textSegment.metadata().put("author","rpj");
        //内容转为向量数据
        Response<Embedding> embeddingResponse = embeddingModel.embed(textSegment);
        //从向量对象中获取向量内容
        Embedding embedding = embeddingResponse.content();
        //将向量数据添加到向量数据库中
        String result = embeddingStore.add(embedding, textSegment);
        System.out.println(result);
        return result;
    }



    @GetMapping(value = "/embedding/query1")
    public void query1(){
        Embedding queryEmbedding = embeddingModel.embed("咏鸡说的是什么").content();
        //作用和mybatis的wapper作用相似，都是包装查询条件
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)//查询条件
                .maxResults(1)//最大搜索数量
                .build();
        //查询引擎 表示一个向量存储库，获取查询结果  --检索增强生成
        EmbeddingSearchResult<TextSegment> searchResult =
                embeddingStore.search(embeddingSearchRequest);
        //只取第一条查询结果，转为文本输出
        System.out.println(searchResult.matches().get(0).embedded().text());

    }


    @GetMapping(value = "/embedding/query2")
    public void query2(){
        Embedding embedding = embeddingModel.embed("咏鸡").content();

        //查询咏鹅作者是rpj2
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(embedding)
                .filter(metadataKey("author").isEqualTo("rpj2"))
                .maxResults(1)
                .build();
        //查询向量数据库，获取结果
        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);
        //只取第一条结果，转为文本
        System.out.println(searchResult.matches().get(0).embedded().text());

    }

}
